{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SageMaker V3 Train-to-Inference E2E Example\n",
    "\n",
    "This notebook demonstrates the complete end-to-end workflow from training a custom PyTorch model to deploying it for inference using SageMaker V3."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prerequisites\n",
    "Note: Ensure you have sagemaker and ipywidgets installed in your environment. The ipywidgets package is required to monitor endpoint deployment progress in Jupyter notebooks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import required libraries\n",
    "import json\n",
    "import uuid\n",
    "import tempfile\n",
    "import os\n",
    "import boto3\n",
    "\n",
    "from sagemaker.serve.model_builder import ModelBuilder\n",
    "from sagemaker.serve.builder.schema_builder import SchemaBuilder\n",
    "from sagemaker.serve.utils.types import ModelServer\n",
    "from sagemaker.serve.spec.inference_spec import InferenceSpec\n",
    "from sagemaker.train.model_trainer import ModelTrainer\n",
    "from sagemaker.train.configs import SourceCode\n",
    "from sagemaker.core.resources import EndpointConfig\n",
    "from sagemaker.core.helper.session_helper import Session"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 1: Configure Training Job\n",
    "\n",
    "Set up a custom PyTorch model for training. We'll create a simple neural network for demonstration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Configuration for training\n",
    "MODEL_NAME_PREFIX = \"train-inf-v3-example-model\"\n",
    "ENDPOINT_NAME_PREFIX = \"train-inf-v3-example-endpoint\"\n",
    "TRAINING_JOB_PREFIX = \"e2e-v3-pytorch\"\n",
    "\n",
    "# AWS Configuration\n",
    "AWS_REGION = \"us-west-2\"\n",
    "PYTORCH_TRAINING_IMAGE = \"763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:1.13.1-cpu-py39\"\n",
    "\n",
    "# Generate unique identifiers\n",
    "unique_id = str(uuid.uuid4())[:8]\n",
    "training_job_name = f\"{TRAINING_JOB_PREFIX}-{unique_id}\"\n",
    "model_name = f\"{MODEL_NAME_PREFIX}-{unique_id}\"\n",
    "endpoint_name = f\"{ENDPOINT_NAME_PREFIX}-{unique_id}\"\n",
    "\n",
    "print(f\"Training job name: {training_job_name}\")\n",
    "print(f\"Model name: {model_name}\")\n",
    "print(f\"Endpoint name: {endpoint_name}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 2: Create Training Code\n",
    "\n",
    "Create a simple PyTorch training script and requirements file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_pytorch_training_code():\n",
    "    \"\"\"Create PyTorch training script.\"\"\"\n",
    "    temp_dir = tempfile.mkdtemp()\n",
    "    \n",
    "    train_script = '''import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import DataLoader, TensorDataset\n",
    "import os\n",
    "\n",
    "class SimpleModel(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.linear = nn.Linear(4, 2)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        return torch.softmax(self.linear(x), dim=1)\n",
    "\n",
    "def train():\n",
    "    model = SimpleModel()\n",
    "    optimizer = optim.Adam(model.parameters(), lr=0.01)\n",
    "    criterion = nn.CrossEntropyLoss()\n",
    "    \n",
    "    # Synthetic data\n",
    "    X = torch.randn(100, 4)\n",
    "    y = torch.randint(0, 2, (100,))\n",
    "    dataset = TensorDataset(X, y)\n",
    "    dataloader = DataLoader(dataset, batch_size=32)\n",
    "    \n",
    "    # Train for 1 epoch\n",
    "    model.train()\n",
    "    for batch_x, batch_y in dataloader:\n",
    "        optimizer.zero_grad()\n",
    "        outputs = model(batch_x)\n",
    "        loss = criterion(outputs, batch_y)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "    \n",
    "    # Save model for TorchServe\n",
    "    model.eval()\n",
    "    traced_model = torch.jit.trace(model, torch.randn(1, 4))\n",
    "    \n",
    "    model_dir = os.environ.get('SM_MODEL_DIR', '/opt/ml/model')\n",
    "    os.makedirs(model_dir, exist_ok=True)\n",
    "    torch.jit.save(traced_model, os.path.join(model_dir, 'model.pth'))\n",
    "    \n",
    "    print(\"Training completed and model saved!\")\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    train()\n",
    "'''\n",
    "    \n",
    "    with open(os.path.join(temp_dir, 'train.py'), 'w') as f:\n",
    "        f.write(train_script)\n",
    "    \n",
    "    with open(os.path.join(temp_dir, 'requirements.txt'), 'w') as f:\n",
    "        f.write('torch>=1.13.0,<2.0.0\\n')\n",
    "    \n",
    "    return temp_dir\n",
    "\n",
    "# Create training code\n",
    "training_code_dir = create_pytorch_training_code()\n",
    "print(f\"Training code created in: {training_code_dir}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 3: Create ModelTrainer and Start Training\n",
    "\n",
    "Set up the ModelTrainer with custom PyTorch code and launch the training job."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create SageMaker session\n",
    "boto_session = boto3.Session(region_name=AWS_REGION)\n",
    "sagemaker_session = Session(boto_session=boto_session)\n",
    "\n",
    "# Create ModelTrainer with custom code\n",
    "model_trainer = ModelTrainer(\n",
    "    sagemaker_session=sagemaker_session,\n",
    "    training_image=PYTORCH_TRAINING_IMAGE,\n",
    "    source_code=SourceCode(\n",
    "        source_dir=training_code_dir,\n",
    "        entry_script=\"train.py\",\n",
    "        requirements=\"requirements.txt\",\n",
    "    ),\n",
    "    base_job_name=training_job_name\n",
    ")\n",
    "\n",
    "# Start training job\n",
    "print(f\"Starting training job: {training_job_name}\")\n",
    "print(\"Note: This will take a few minutes to complete.\")\n",
    "\n",
    "model_trainer.train()\n",
    "print(\"Model Training Completed!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 4: Create Schema Builder and Inference Spec\n",
    "\n",
    "Set up the schema and inference specification for the trained model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create schema builder for tensor-based models\n",
    "def create_schema_builder():\n",
    "    \"\"\"Create schema builder for tensor-based models.\"\"\"\n",
    "    sample_input = [[0.1, 0.2, 0.3, 0.4]]\n",
    "    sample_output = [[0.8, 0.2]]\n",
    "    return SchemaBuilder(sample_input, sample_output)\n",
    "\n",
    "# Create inference specification\n",
    "class SimpleInferenceSpec(InferenceSpec):\n",
    "    def load(self, model_dir):\n",
    "        import torch\n",
    "        return torch.jit.load(f\"{model_dir}/model.pth\")\n",
    "    \n",
    "    def invoke(self, input_object, model):\n",
    "        import torch\n",
    "        return model(torch.tensor(input_object)).tolist()\n",
    "\n",
    "schema_builder = create_schema_builder()\n",
    "print(\"Schema builder and inference spec created successfully!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 5: Create ModelBuilder and Build Model\n",
    "\n",
    "Create the ModelBuilder with the trained model and build it for deployment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create ModelBuilder with trained model\n",
    "model_builder = ModelBuilder(\n",
    "    model=model_trainer,\n",
    "    schema_builder=schema_builder,\n",
    "    model_server=ModelServer.TORCHSERVE,\n",
    "    inference_spec=SimpleInferenceSpec(),\n",
    "    image_uri=PYTORCH_TRAINING_IMAGE.replace(\"training\", \"inference\"),\n",
    "    dependencies={\"auto\": False},\n",
    ")\n",
    "\n",
    "# Build the trained model\n",
    "core_model = model_builder.build(model_name=model_name, region=AWS_REGION)\n",
    "print(f\"Model Successfully Created: {core_model.model_name}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 6: Deploy the Trained Model\n",
    "\n",
    "Deploy the trained model to a SageMaker endpoint."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Deploy the trained model\n",
    "core_endpoint = model_builder.deploy(\n",
    "    endpoint_name=endpoint_name,\n",
    "    initial_instance_count=1\n",
    ")\n",
    "print(f\"Endpoint Successfully Created: {core_endpoint.endpoint_name}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 7: Test the Trained Model\n",
    "\n",
    "Test the deployed trained model with sample tensor inputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test the trained model with tensor input\n",
    "test_data = [[0.1, 0.2, 0.3, 0.4]]\n",
    "\n",
    "result = core_endpoint.invoke(\n",
    "    body=json.dumps(test_data),\n",
    "    content_type=\"application/json\"\n",
    ")\n",
    "\n",
    "# Decode and display the result\n",
    "prediction = json.loads(result.body.read().decode('utf-8'))\n",
    "print(f\"Result of invoking endpoint: {prediction}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test with different tensor inputs\n",
    "test_inputs = [\n",
    "    [[0.5, 0.3, 0.2, 0.1]],\n",
    "    [[0.9, 0.1, 0.8, 0.2]],\n",
    "    [[0.2, 0.7, 0.4, 0.6]]\n",
    "]\n",
    "\n",
    "for i, test_input in enumerate(test_inputs, 1):\n",
    "    result = core_endpoint.invoke(\n",
    "        body=json.dumps(test_input),\n",
    "        content_type=\"application/json\"\n",
    "    )\n",
    "    \n",
    "    prediction = json.loads(result.body.read().decode('utf-8'))\n",
    "    print(f\"Test {i} - Input {test_input}: {prediction}\")\n",
    "    print(\"-\" * 50)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 8: Clean Up Resources\n",
    "\n",
    "Clean up all created resources including the trained model and endpoint."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Clean up resources\n",
    "core_endpoint_config = EndpointConfig.get(endpoint_config_name=core_endpoint.endpoint_name)\n",
    "\n",
    "# Delete in the correct order\n",
    "core_model.delete()\n",
    "core_endpoint.delete()\n",
    "core_endpoint_config.delete()\n",
    "\n",
    "print(\"Model and Endpoint Successfully Deleted!\")\n",
    "print(f\"Note: Training job artifacts remain in S3 for reference.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "This notebook demonstrated the complete E2E workflow:\n",
    "1. Creating custom PyTorch training code\n",
    "2. Configuring a ModelTrainer with custom source code\n",
    "3. Running a training job on SageMaker\n",
    "4. Building a ModelBuilder from training artifacts\n",
    "5. Deploying the trained model to an endpoint\n",
    "6. Testing the trained model with tensor inputs\n",
    "7. Proper cleanup of inference resources\n",
    "\n",
    "## Key Benefits of E2E Training:\n",
    "- **Custom training**: Full control over PyTorch training process\n",
    "- **Seamless workflow**: Train → Build → Deploy in one pipeline\n",
    "- **Artifact management**: Automatic handling of training outputs\n",
    "- **TorchServe integration**: Easy deployment with TorchServe\n",
    "- **Production ready**: Trained models ready for immediate deployment\n",
    "\n",
    "The V3 ModelBuilder makes it easy to go from custom training to production inference with minimal code!"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv-test",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
